import scipy
import copy
import random
from helper_codes.general_stimulus import general_stimulus, GeneralStim

def competitive_moving_dot_settings():
    '''
    This function tells you what settings are needed and in what order for this stimulus.
    Inputs:
        None
    Returns:
        settings        : A list of tupples.  In each tupple is
                            (setting_name, setting_description, default_value).
                          The list is of length num_settings + 1, so that the final
                             value in the list is True/False wether or not the stimulus
                             can be triggered via the parallel port or not.
    '''
    settings = []
    settings.append( ('rim_size',
                     'The DIAMETER of the rim on which dots begin and end movement. (normalized RF units)',
                     1.0 ))
    settings.append( ('dot_sizes',
                     'The sizes of the moving dot. (normalized RF units) [list]',
                     [0.2, 0.5] ))
    settings.append( ('duration',
                     'The amount of time the stimulus will move. (seconds)',
                     3.0 ))
    settings.append( ('wait_time',
                     'The amount of time after stimulus movement. (seconds)',
                     1.0 ))
    settings.append( ('angle',
                     'The angle describing where the dot will start moving from. (degrees)',
                     90 ))
    settings.append( ('shade',
                     '1.0 for black dots, -1.0 for white dots.',
                     1.0 ))
    settings.append( ('mirror_pos',
                     'The position of a "mirrored" stimulus relative to the RF. (normalized RF units) [list]',
                     [0.5,0.0] ))
    settings.append( ('delay',
                     'Time between showing the mirror and the regular stimulus. (seconds)',
                     1.0 ))
    settings.append( ('win_size',
                     'The dimensions of the stimulus delivery window. (pixels) [list]',
                     [1440,900] ))
    settings.append( ('refresh_rate',
                     'The number of screen refreshes per second. (Hz)',
                     60 ))
    settings.append( ('trials',
                     'The number of trials run pseudorandomly.',
                     5 ))
    settings.append( ('random_seed',
                     'The seed for the random number generator',
                     40 ))
    settings.append(True) # this stimulus can be triggered
    return settings

class Stim(GeneralStim):
    """
    This is a basic stimulus object, it must have:
        .setup(window) method to set itself up to be drawn
        .draw() method to draw itself to the display
        .evolve() method to set itself up for the next draw
        .exhausted variable set to True after stim is done.
    """
    def __init__( self, rim_size,  dot_sizes, duration, 
                  wait_time, angle, shade, mirror_pos, delay, 
                  random_seed, refresh_rate, trials, preview, save_frames):
        self.rim_size = rim_size
        # this is legacy support for single dot size specifications
        if isinstance(dot_sizes, list): 
            self.dot_sizes = dot_sizes
        else:
            self.dot_sizes = [dot_sizes, dot_sizes]
        self.duration = duration
        self.wait_time = wait_time
        self.angle = angle
        self.shade = shade
        self.mirror_pos = scipy.array(mirror_pos)
        self.delay = delay
        self.refresh_rate = refresh_rate
        self.preview = preview
        self.save_frames = save_frames
        self.trials = trials
        self.random_seed = random_seed

        self.exhausted = False
        self.frame_no = 1
        self.segment_no = 0
        self.trial_no = 1

    def setup(self, window):
        from psychopy import visual
        random.seed(self.random_seed)

        def find_positions(rim_size, angle, num_frames):
            """
            Returns the positions of point moving from one side of a rim to 
                another side of the rim.
            """
            # find the x,y coordinates of the starting point
            start_x = -rim_size/2.0 * scipy.cos(angle*scipy.pi/180.0)
            start_y = -rim_size/2.0 * scipy.sin(angle*scipy.pi/180.0)
            
            speed = rim_size/(num_frames*1.0)

            dx = speed * scipy.cos(angle*scipy.pi/180.0)
            dy = speed * scipy.sin(angle*scipy.pi/180.0)

            positions = [[start_x + i*dx, start_y + i*dy] for i in xrange(num_frames)]

            return positions

        # determine the number of frames per angle.
        num_frames   = int(self.duration*self.refresh_rate)
        wait_frames  = int(self.wait_time*self.refresh_rate)
        delay_frames = int(abs(self.delay)*self.refresh_rate)
        self.num_frames = num_frames
        self.wait_frames = wait_frames
        self.delay_frames = delay_frames

        # determine segment order randomly for each trial.
        if delay_frames:
            possible_segments = ['1-alone', '2-alone', '1-then-2', '2-then-1', '1-and-2-simultaneously']
        else:
            possible_segments = ['1-alone', '2-alone', '1-and-2-simultaneously']

        # we'll create positions and mirror positions for each possible segment.
        target_positions = {}
        mirror_positions = {}

        positions = find_positions(self.rim_size, self.angle, self.num_frames)
        first_pos = positions[0]
        last_pos   = positions[-1]
        positions += [last_pos for i in range(wait_frames)]

        # 1-alone
        target_positions['1-alone'] = copy.deepcopy(positions)
        mirror_positions['1-alone'] = [[1e6,1e6] for i in range(num_frames+wait_frames)] # very far away
        # 2-alone
        target_positions['2-alone'] = [[1e6,1e6] for i in range(num_frames+wait_frames)] # very far away
        mirror_positions['2-alone'] = copy.deepcopy(positions)
        # 1-then-2
        target_positions['1-then-2']  = copy.deepcopy(positions)
        target_positions['1-then-2'] += [last_pos  for i in range(delay_frames)] # add delay frames
        mirror_positions['1-then-2']  = [first_pos for i in range(delay_frames)]
        mirror_positions['1-then-2'] += copy.deepcopy(positions) 
        # 2-then-1
        target_positions['2-then-1'] = copy.deepcopy(mirror_positions['1-then-2'])
        mirror_positions['2-then-1'] = copy.deepcopy(target_positions['1-then-2'])
        # 1-and-2-simultaneously
        target_positions['1-and-2-simultaneously'] = copy.deepcopy(positions)
        mirror_positions['1-and-2-simultaneously'] = copy.deepcopy(positions)

        for key in target_positions.keys():
            target_positions[key] = scipy.array(target_positions[key])
        for key in mirror_positions.keys():
            mirror_positions[key] = scipy.array(mirror_positions[key]) + self.mirror_pos
        self.target_positions = target_positions
        self.mirror_positions = mirror_positions

        self.segment_order = {}
        for trial_no in range(1,self.trials+1):
            candidate_segments = copy.deepcopy(possible_segments)
            segments = []
            for i in xrange(len(possible_segments)):
                index = random.randint(0,len(candidate_segments)-1)
                segments.append(copy.deepcopy(candidate_segments[index]))
                del(candidate_segments[index])
            self.segment_order[trial_no] = copy.deepcopy(segments)
            print "Trial #%d: %s" % (trial_no, str(self.segment_order[trial_no]))

        # construct the patch_info for the bar
        texture = scipy.ones((2,2),dtype=scipy.float64)*-1.0*self.shade
        self.dot = visual.PatchStim( window, 
                                     tex=texture, 
                                     pos = target_positions[self.segment_order[1][0]][0],
                                     size=(self.dot_sizes[0], self.dot_sizes[0]), 
                                     mask = 'circle',
                                     units='norm')
        self.mirror_dot = visual.PatchStim( window, 
                                     tex=texture, 
                                     pos = mirror_positions[self.segment_order[1][0]][0],
                                     size=(self.dot_sizes[1], self.dot_sizes[1]), 
                                     mask = 'circle',
                                     units='norm')
        self.textures = [self.dot, self.mirror_dot]
        self.trial_no = 1


    def draw(self):
        self.dot.draw()
        self.mirror_dot.draw()

    def evolve(self):
        try:
            # find out what segment we should be showing.
            this_segment = self.segment_order[self.trial_no][self.segment_no]
        except KeyError:
            # end unless in preview mode
            if (not self.preview) or self.save_frames:
                self.exhausted = True
                return
            else:
                # reset everything
                self.frame_no = 0
                self.segment_no = 0
                self.trial_no = 1
                self.evolve()
                return
        except IndexError:
            self.frame_no = 0
            self.segment_no = 0
            self.trial_no += 1
            self.evolve()
            return
            
        try:
            self.dot.setPos(       self.target_positions[this_segment][self.frame_no])
            self.mirror_dot.setPos(self.mirror_positions[this_segment][self.frame_no])
            self.frame_no += 1
        except IndexError:
            self.frame_no    = 0
            self.segment_no += 1
            self.evolve()
            return

def competitive_moving_dot_runtime(rim_size, dot_sizes, duration, wait_time, 
                      angle, shade, mirror_pos, delay, win_size, 
                      refresh_rate, trials, random_seed):
    """
    Calculate and return the number of seconds that this stimulus will need to run, after being triggered.
    """
    # for each 'trial' we will go through each permutation of stimuli order, including simulaneous.
    #    we have 1-alone, 2-alone, 1-then-2, 2-then-1, 1-and-2-simultaneous
    #    if delay=0 then discard cases 1-then-2 and 2-then-1. 
    total_time = 0.0
    if delay:
        #               1    2   1->2    2->1   1<>2
        for tdelay in [0.0, 0.0, delay, -delay, 0.0]:
            total_time +=  (duration + wait_time) + abs(tdelay)
    else:
        #               1    2   1<>2
        for tdelay in [0.0, 0.0, 0.0]:
            total_time +=  (duration + wait_time) + abs(tdelay)
    total_time *= trials 
    return total_time + 0.99 # 0.99 added so rounding is always up.

def competitive_moving_dot(rim_size, dot_sizes, duration, wait_time, 
                      angle, shade, mirror_pos, delay, win_size, 
                      refresh_rate, trials, random_seed, 
                         window=None,
                         eye_cup=None, 
                         rf=None, 
                         save_frames=False,
                         save_frames_factor=1.0,
                         save_frames_size=100,
                         display_saved_frames=False,
                         preview=True):
    """
    This stimulus shows a dot moving radially through the
        receptive field.  A second dot is displayed offset from
        the center of the RF by mirror_pos and with a temporal delay.
    If trials is not zero, then a set of pseudorandomly determined segments
        is constructed.
    Inputs:
        rim_size        : the radius of the rim where dots begin and end in normalized units.
        dot_sizes       : the size of the DIAMETER of the moving dot in normalized units.
        duration        : the time where the dot will be moving
        wait_time       : the time after movement, before screen returns to previous state.
        angles          : angle through which the dot moves.
        mirror_pos      : the position of a "mirror" of the regular stimulus in normalized units.
        delay           : delay between regular and mirror display, negative values supported.
        win_size        : dimensions of the screen in pixels, (1440,900) is default
        refresh_rate    : refresh rate in Hz of the video, 60.0 is default.
        random_seed     : seed which sets the random number generator
        trials          : number of trials to run
        --- kwargs --
        window          : a psychopy.visual Window object, if not passed it will be constructed.
                        :    if passed it will not be closed (a blank screen will be drawn)
        preview         : If false the program will wait for a pulse trigger, otherwise it won't.
        save_frames_factor  : if None, then no frames are saved, otherwise it is a 
                                scaling factor so we can scale down the frames and
                                output a frames.gif (an animated gif file).
        eye_cup         : an Eye_cup object, see its docstring for more info.
        rf              : RF object, to set up new coordinate system.
    Returns:
        times     : The times when frames were shown, with zero being the first frame time.
    """
    if trials <= 0:
        trials = 1
    stim_list = [ Stim( rim_size, dot_sizes, duration, 
                        wait_time, angle, shade, mirror_pos, delay, 
                        random_seed, refresh_rate, trials, preview, save_frames) ]
    return_values = general_stimulus( stim_list, shade, win_size, 
                                      window=window,
                                      eye_cup=eye_cup, 
                                      rf=rf, 
                                      save_frames=save_frames,
                                      save_frames_factor=save_frames_factor,
                                      save_frames_size=save_frames_size,
                                      display_saved_frames=display_saved_frames,
                                      preview=preview)
    return return_values


if __name__ == '__main__':
    # ---------------------------------
    # -- Set the stimulus parameters --
    # ---------------------------------
    factor      = 0.8
    rim_size    = 1.0         # diameter of the rim where dots start
    dot_sizes   = [0.2, 0.6]         # diameter of dot in RF units
    duration    = 2.0         # in seconds
    wait_time   = 1.0         # in seconds
    angle       = 90          # the angle the dot will move through
    random_seed = 40
    trials      = 2
    shade       = -1.0        # 1.0 for white, -1.0 for black background
    mirror_pos  = [0.5,0.0]
    delay       = 0.5

    # display parameters
    win_size = (1440,900)  # in pixels
    refresh_rate = 60.0    # in Hz
    return_values = competitive_moving_dot(rim_size, dot_sizes, duration, wait_time, 
                      angle, shade, mirror_pos, delay, win_size, refresh_rate, 
                      random_seed, trials, 
                      save_frames = False,
                      preview=True)
